
#!/bin/bash
set -e  # 在任何命令失败时退出

# 设置字符编码为UTF-8
export LANG=zh_CN.UTF-8

echo "========================================"
echo "           LoRA 训练脚本"
echo "========================================"
echo
WORKSPACE="/workspace"



# 2. 激活虚拟环境
source "$WORKSPACE/AI/.train/bin/activate"

# 设置Python路径 (请根据你的Linux环境修改)
PYTHON_PATH="/workspace/AI/.train/bin/python3"
WORKSPACE="$(dirname "$0")"

# 设置脚本路径
SCRIPT_PATH="./nodes/sd_train/train_network.py"
echo "正在调用脚本: $SCRIPT_PATH"


# 执行模式选择
choose_mode() {
    echo "========================================"
    echo "           选择执行模式"
    echo "========================================"
    echo "1. thread     - 线程模式 (推荐)"
    echo "   在当前进程中使用线程执行训练"
    echo "   资源共享，速度较快，适合大多数情况"
    echo
    echo "2. subprocess - 子进程模式"
    echo "   创建独立子进程执行训练"
    echo "   资源隔离，更稳定，适合长时间训练"
    echo
    echo "========================================"

    read -p "请输入选择 (1 或 2，直接回车默认选择1): " mode_choice
    
    if [ -z "$mode_choice" ]; then
        mode_choice=1
    fi
    
    if [ "$mode_choice" == "1" ]; then
        EXECUTION_MODE="thread"
        echo "已选择: 线程模式 (thread)"
    elif [ "$mode_choice" == "2" ]; then
        EXECUTION_MODE="subprocess"
        echo "已选择: 子进程模式 (subprocess)"
    else
        echo "无效选择，请输入 1 或 2"
        choose_mode
    fi
}

choose_mode

echo
# 设置路径 (请根据你的Linux环境修改这些路径)
MODEL_PATH="/workspace/AI/ComfyUI/models/checkpoints/y.safetensors"
DATA_DIR="/workspace/training_data/images"
OUTPUT_DIR="/workspace/training_data/loraout"

mkdir -p "$OUTPUT_DIR" "$OUTPUT_DIR/logs"
# 验证训练数据目录
if [ ! -d "$DATA_DIR" ]; then
    echo "错误: 训练数据目录不存在: $DATA_DIR"
    exit 1
fi

# 验证模型路径
if [ ! -f "$MODEL_PATH" ]; then
    echo "错误: 模型文件不存在: $MODEL_PATH"
    exit 1
fi

OUTPUT_NAME="my_lora"
NETWORK_DIM=32
NETWORK_ALPHA=32
BATCH_SIZE=1
LR=0.0001
STEPS=10
RESOLUTION=512

# 设置 PYTHONPATH
SCRIPT_DIR="$(dirname "$0")"
export PYTHONPATH="$SCRIPT_DIR:$SCRIPT_DIR/sd_scripts:$PYTHONPATH"
echo "PYTHONPATH设置为: $PYTHONPATH"
echo
echo "========================================"
echo "开始训练，执行模式: $EXECUTION_MODE"
echo "========================================"

echo 设置huggingface镜像
export HF_ENDPOINT=https://hf-mirror.com

# 执行命令并记录日志
{
    $PYTHON_PATH "$SCRIPT_PATH" \
        --pretrained_model_name_or_path "$MODEL_PATH" \
        --train_data_dir "$DATA_DIR" \
        --output_dir "$OUTPUT_DIR" \
        --output_name "$OUTPUT_NAME" \
        --tokenizer_cache_dir "/workspace/AI/tokenizer" \
        --network_dim $NETWORK_DIM \
        --network_alpha $NETWORK_ALPHA \
        --train_batch_size $BATCH_SIZE \
        --learning_rate $LR \
        --max_train_steps $STEPS \
        --save_model_as safetensors \
        --mixed_precision fp16 \
        --save_every_n_steps 500 \
        --seed 42 \
        --resolution "$RESOLUTION,$RESOLUTION" \
        --caption_extension .txt \
        --network_module networks.lora \
        --bucket_reso_steps=64 \
        --min_bucket_reso=320 \
        --max_bucket_reso=1280 \
        --enable_bucket \
        --logging_dir "$OUTPUT_DIR/logs" \
        --cache_latents \
        --xformers \
        --gradient_checkpointing \
        --cache_latents_to_disk \
        --gradient_accumulation_steps 1

} 2>&1 | tee "$OUTPUT_DIR/training.log"


              
# 检查执行结果
if [ $? -ne 0 ]; then
    echo "训练失败，错误代码: $?"
    exit $?
fi

echo "训练完成!"
read -p "按回车键退出..."
